import numpy as np
import cv2
import math
import imageio
import os
import datetime
import pandas as pd
from lib import Heading

# 配置参数
class Config:
    def __init__(self):
        self.v_max = 1.5
        self.v_min = 0.0
        self.omega_max = np.pi / 2
        self.omega_min = -np.pi / 2
        self.a_max = 1.5
        self.alpha_max = np.pi / 1
        self.dt = 0.1  # 提高采样频率（原为0.2）
        self.predict_time = 1.5
        self.v_res = 0.02  # 更细的速度采样
        self.omega_res = np.pi / 180  # 更细的角速度采样
        self.w_heading = 0.3  # 目标距离权重
        self.w_clear = 0.4  # 障碍物距离权重
        self.w_velocity = 0.3  # 速度权重
        self.robot_radius = 0.1
        self.obstacle_radius = 0.1
        self.obstacle_radius_safe = 0.5

config = Config()
# 创建输出文件夹
now_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
out_dir = f"DWA_Output_{now_str}"
os.makedirs(out_dir, exist_ok=True)

# 初始状态和目标
x = np.array([0, 0, 0.785, 1.5, 0])  # [x, y, yaw, v, omega]
goal = np.array([10, 10])

# 障碍物 [x, y, vx, vy]
obstacles = np.array([
    [4.5, 5, 0.1, 0],
    [6, 8, -0.05, -0.05],
    [7, 5, 0, 0.1]
])

trajectory = [x.copy()]
candidate_trajs = []
best_traj_global = None

# 可视化画布
scale = 50
canvas_size = (600, 600)
frames = []  # 存储GIF帧

def draw(canvas, robot, goal, obstacles, traj, candidate_trajs, best_traj):
    canvas[:] = 255
    for obs in obstacles:
        obs_px = (int(obs[0]*scale), int(obs[1]*scale))
        r = int((config.obstacle_radius)*scale)
        r_safe = int((config.obstacle_radius_safe)*scale)
        cv2.circle(canvas, obs_px, r, (0,0,0), -1)
        cv2.circle(canvas, obs_px, r_safe, (0,0,0), 2)
    for cand in candidate_trajs:
        for i in range(1, len(cand)):
            p1 = (int(cand[i-1][0]*scale), int(cand[i-1][1]*scale))
            p2 = (int(cand[i][0]*scale), int(cand[i][1]*scale))
            cv2.line(canvas, p1, p2, (200, 200, 200), 1)
    if best_traj is not None:
        for i in range(1, len(best_traj)):
            p1 = (int(best_traj[i-1][0]*scale), int(best_traj[i-1][1]*scale))
            p2 = (int(best_traj[i][0]*scale), int(best_traj[i][1]*scale))
            cv2.line(canvas, p1, p2, (0, 0, 255), 2)
    for i in range(1, len(traj)):
        p1 = (int(traj[i-1][0]*scale), int(traj[i-1][1]*scale))
        p2 = (int(traj[i][0]*scale), int(traj[i][1]*scale))
        cv2.line(canvas, p1, p2, (255, 0, 0), 2)
    cv2.circle(canvas, (int(goal[0]*scale), int(goal[1]*scale)), 5, (0,0,255), -1)
    robot_px = (int(robot[0]*scale), int(robot[1]*scale))
    cv2.circle(canvas, robot_px, int(config.robot_radius*scale), (255, 0, 255), 2)

# 运动学模型
def motion(x, u, dt):
    x_new = x.copy()
    x_new[0] += u[0] * np.cos(x[2]) * dt
    x_new[1] += u[0] * np.sin(x[2]) * dt
    x_new[2] += u[1] * dt
    x_new[3] = u[0]
    x_new[4] = u[1]
    return x_new

def calc_dynamic_window(x):
    Vs = [config.v_min, config.v_max, config.omega_min, config.omega_max]
    Vd = [x[3]-config.a_max*config.dt, x[3]+config.a_max*config.dt,
          x[4]-config.alpha_max*config.dt, x[4]+config.alpha_max*config.dt]
    return [max(Vs[0], Vd[0]), min(Vs[1], Vd[1]), max(Vs[2], Vd[2]), min(Vs[3], Vd[3])]

def predict_trajectory(x_init, u):
    x = x_init.copy()
    traj = [x[:3].copy()]
    time = 0
    while time <= config.predict_time:
        x = motion(x, u, config.dt)
        traj.append(x[:3].copy())
        time += config.dt
    return np.array(traj)

def calc_heading_score(traj, goal):
    dx = goal[0] - traj[-1,0]
    dy = goal[1] - traj[-1,1]
    target_angle = np.arctan2(dy, dx)
    yaw = traj[-1,2]
    error_angle = np.arctan2(np.sin(target_angle - yaw), np.cos(target_angle - yaw))
    return np.cos(error_angle)

def calc_clearance_score(traj, obstacles):
    min_dist = float('inf')
    time = 0
    for point in traj:
        for obs in obstacles:
            obs_pos = obs[0:2] + obs[2:4] * time
            dist = np.linalg.norm(point[:2] - obs_pos)
            clearance = dist - config.robot_radius - config.obstacle_radius
            if clearance < min_dist:
                min_dist = clearance
        time += config.dt
    min_dist = max(min_dist, 0.01)
    return np.exp(-1 / min_dist)

def dwa_control(x, goal, obstacles):
    dw = calc_dynamic_window(x)
    max_score = -float('inf')
    best_u = [0, 0]
    best_traj = None
    local_candidates = []

    for v in np.arange(dw[0], dw[1], config.v_res):
        for omega in np.arange(dw[2], dw[3], config.omega_res):
            traj = predict_trajectory(x, [v, omega])
            local_candidates.append(traj)
            heading = calc_heading_score(traj, goal)
            clearance = calc_clearance_score(traj, obstacles)
            score = config.w_heading*heading + config.w_clear*clearance + config.w_velocity*v
            if score > max_score:
                max_score = score
                best_u = [v, omega]
                best_traj = traj

    return best_u, best_traj, local_candidates

# 主循环
canvas = np.ones((canvas_size[1], canvas_size[0], 3), dtype=np.uint8) * 255
heading_records = []
speed_records = []
path_length = 0.0  # 初始化路径长度

while np.linalg.norm(x[:2] - goal) > 0.5:
    obstacles[:,0:2] += obstacles[:,2:4] * config.dt
    prev_position = x[:2].copy()
    u, best_traj, candidate_trajs = dwa_control(x, goal, obstacles)
    x = motion(x, u, config.dt)
    trajectory.append(x.copy())
    current_position = x[:2]
    path_length += np.linalg.norm(current_position - prev_position)
    heading_records.append(x[2]*57.3)  # 转为角度进行记录
    speed_records.append(x[3])
    draw(canvas, x, goal, obstacles, trajectory, candidate_trajs, best_traj)
    frame = cv2.cvtColor(canvas.copy(), cv2.COLOR_BGR2RGB)
    frames.append(frame)
    cv2.imshow('DWA Simulation', canvas)
    key = cv2.waitKey(1)
    if key == 27:
        break

# 导出GIF
gif_path = os.path.join(out_dir, "trajectory.gif")
imageio.mimsave(gif_path, frames, duration=0.05)


evaluator = Heading.PathSmoothnessEvaluator(heading_records, unit='deg')
result = evaluator.evaluate()
print(f"路径总长度: {path_length}")
print(f"累计转角: {result['total_turning_angle']:.2f}°")
print(f"最大单次转角: {result['max_turning_angle']:.2f}°")

# 导出Excel
timestamps = [i for i in range(len(heading_records))]
df_records = pd.DataFrame({
    'Time': timestamps,
    'Speed': speed_records,
    'Yaw': heading_records

})

df_records.to_excel(os.path.join(out_dir, "velocity_heading.xlsx"), index=False)


cv2.destroyAllWindows()
print(f"✅ 动图和航速航向数据已保存至文件夹：{out_dir}")